using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Threading;

using Supremacy.Annotations;

namespace SupremacyEditor
{
    public class ObservableKeyedCollection<TKey, TValue> : ObservableCollection<TValue>
        where TValue : class
    {
        private readonly ReaderWriterLockSlim _lock;
        private readonly Func<TValue, TKey> _keyResolver;
        private readonly Dictionary<TKey, int> _keyIndexLookup;

        public TValue this[[NotNull] TKey key]
        {
            get
            {
                _lock.EnterReadLock();
                try
                {
                    int index;
                    if (_keyIndexLookup.TryGetValue(key, out index))
                        return base[index];
                    return null;
                }
                finally
                {
                    _lock.ExitReadLock();
                }
            }
        }

        public ObservableKeyedCollection([NotNull] Func<TValue, TKey> keyResolver)
        {
            if (keyResolver == null)
                throw new ArgumentNullException("keyResolver");
            _lock = new ReaderWriterLockSlim(LockRecursionPolicy.SupportsRecursion);
            _keyResolver = keyResolver;
            _keyIndexLookup = new Dictionary<TKey, int>();
        }

        public void AddRange([NotNull] IEnumerable<TValue> items)
        {
            if (items == null)
                throw new ArgumentNullException("items");
            foreach (var item in items)
                Add(item);
        }

        protected override void ClearItems()
        {
            _lock.EnterWriteLock();
            try
            {
                base.ClearItems();
                _keyIndexLookup.Clear();
            }
            finally
            {
                _lock.ExitWriteLock();
            }
        }

        protected override void InsertItem(int index, TValue item)
        {
            _lock.EnterWriteLock();
            try
            {
                base.InsertItem(index, item);
                _keyIndexLookup[_keyResolver(item)] = index;
            }
            finally
            {
                _lock.ExitWriteLock();
            }
        }

        protected override void MoveItem(int oldIndex, int newIndex)
        {
            _lock.EnterWriteLock();
            try
            {
                var item = base[oldIndex];
                base.MoveItem(oldIndex, newIndex);
                _keyIndexLookup[_keyResolver(item)] = newIndex;
                if (newIndex < oldIndex)
                {
                    for (int i = newIndex + 1; i <= oldIndex; i++)
                    {
                        item = base[i];
                        _keyIndexLookup[_keyResolver(item)] = i;
                    }
                }
                else
                {
                    for (int i = oldIndex; i < newIndex; i++)
                    {
                        item = base[i];
                        _keyIndexLookup[_keyResolver(item)] = i;
                    }
                }
            }
            finally
            {
                _lock.ExitWriteLock();
            }
        }

        protected override void SetItem(int index, TValue item)
        {
            _lock.EnterWriteLock();
            try
            {
                base.SetItem(index, item);
                _keyIndexLookup[_keyResolver(item)] = index;
            }
            finally
            {
                _lock.ExitWriteLock();
            }
        }

        protected override void RemoveItem(int index)
        {
            _lock.EnterWriteLock();
            try
            {
                var item = base[index];
                base.RemoveItem(index);
                _keyIndexLookup.Remove(_keyResolver(item));
            }
            finally
            {
                _lock.ExitWriteLock();
            }
        }
    }
}